{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Wavefront sensing with a Pyramid wavefront sensor\n",
    "\n",
    "We will simulate a closed-loop adaptive optics system, based on the the Magellan Adaptive Optics Extreme (MagAO-X) system, that uses an unmodulated pyramid wavefront sensor with a 2k-MEMS DM.\n",
    "\n",
    "We first start by importing the relevant python modules."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from hcipy import *\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# These modules are used for animating some of the graphs in our notebook.\n",
    "from matplotlib import animation, rc\n",
    "from IPython.display import HTML\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We start by defining a few parameters according to the MagAO-X specifications. The Magallen telescope has a diameter of 6.5 meters, and we will use a sensing wavelength of 842nm. A zero magnitude star will have flux of 3.9E10 photons/s."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wavelength_wfs = 842.0E-9\n",
    "telescope_diameter = 6.5\n",
    "zero_magnitude_flux = 3.9E10\n",
    "stellar_magnitude = 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The pyramid wavefront sensor design has a sampling of 56 pixels across the pupil with a distance of 60 pixels between the pupils. The final image size will be 120x120, which is the sampling of OCAM2K camera after 2x2 binning. To get the sampling correct it is the easiest to choose the pupil_grid_diameter as 60/56 larger than the actual telescope_diameter. On the created pupil grids we can evaluate the telescope aperture."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_pupil_pixels = 60\n",
    "pupil_grid_diameter = 60/56 * telescope_diameter\n",
    "pupil_grid = make_pupil_grid(num_pupil_pixels, pupil_grid_diameter)\n",
    "\n",
    "magellan_aperture = evaluate_supersampled(make_magellan_aperture(), pupil_grid, 6)\n",
    "\n",
    "imshow_field(magellan_aperture)\n",
    "plt.xlabel('x position(m)')\n",
    "plt.ylabel('y position(m)')\n",
    "plt.colorbar()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's make our deformable mirror. MagAO-X uses a 2k-MEMS DM of Boston Micromachines. The influence functions of the DM are nearly gaussian. We will therefore make a DM with Gaussian influence functions. There are 50 actuators across the pupil. But for speed purposes we will limit the number of actuators to 10 across the pupil."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_actuators_across_pupil = 10\n",
    "actuator_spacing = telescope_diameter / num_actuators_across_pupil\n",
    "influence_functions = make_gaussian_influence_functions(pupil_grid, num_actuators_across_pupil, actuator_spacing)\n",
    "deformable_mirror = DeformableMirror(influence_functions)\n",
    "num_modes = deformable_mirror.num_actuators\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we are going to make the optics of the pyramid wavefront sensor and the camera. Because the OCAM2K is a very high performance EMCCD we will simulate this detector as a noiseless detector."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pwfs = PyramidWavefrontSensorOptics(pupil_grid, wavelength_0=wavelength_wfs)\n",
    "camera = NoiselessDetector()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are going to use a linear reconstruction algorithm for the wavefront estimation and for that we will need to measure the reference response of a perfect incoming wavefront. To create this we create an unabberated wavefront and propagate it through the pyramid wavefront sensor. Then we will integrate the response with our camera.\n",
    "\n",
    "The final reference will be divided by the total sum to normalize the wavefront sensor response. Doing this consequently for all exposures will make sure that we can use this reference for arbitrary exposure times and photon fluxes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wf = Wavefront(magellan_aperture, wavelength_wfs)\n",
    "wf.total_power = 1\n",
    "\n",
    "camera.integrate(pwfs.forward(wf), 1)\n",
    "\n",
    "image_ref = camera.read_out()\n",
    "image_ref /= image_ref.sum()\n",
    "\n",
    "imshow_field(image_ref)\n",
    "plt.colorbar()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the linear reconstructor we need to now the interaction matrix, which tells us how the pyramid wavefront sensor responds to each actuator of the deformable mirror. This can be build by sequentially applying a positive and negative voltage on a single actuator. The difference between the two gives us the actuator response.\n",
    "\n",
    "We will use the full image of the pyramid wavefront sensor for the reconstruction, so we do not compute the normalized differences between the pupils."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the interaction matrix\n",
    "probe_amp = 0.01 * wavelength_wfs\n",
    "slopes = []\n",
    "\n",
    "wf = Wavefront(magellan_aperture, wavelength_wfs)\n",
    "wf.total_power = 1\n",
    "\n",
    "for ind in range(num_modes):\n",
    "    if ind % 10 == 0:\n",
    "        print(\"Measure response to mode {:d} / {:d}\".format(ind+1, num_modes))\n",
    "    slope = 0\n",
    "\n",
    "    # Probe the phase response\n",
    "    for s in [1, -1]:\n",
    "        amp = np.zeros((num_modes,))\n",
    "        amp[ind] = s * probe_amp\n",
    "        deformable_mirror.actuators = amp\n",
    "\n",
    "        dm_wf = deformable_mirror.forward(wf)\n",
    "        wfs_wf = pwfs.forward(dm_wf)\n",
    "\n",
    "        camera.integrate(wfs_wf, 1)\n",
    "        image = camera.read_out()\n",
    "        image /= np.sum(image)\n",
    "\n",
    "        slope += s * (image-image_ref)/(2 * probe_amp)\n",
    "\n",
    "    slopes.append(slope)\n",
    "\n",
    "slopes = ModeBasis(slopes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The matrix that we build by poking the actuators can be used to transform a DM pattern into the wavefront sensor response. For wavefront reconstruction we want to invert this. We currently have,\n",
    "\n",
    "$$\\vec{S} = A\\vec{\\phi}.$$\n",
    "\n",
    "With $\\vec{S}$ being the response of the wavefront sensor, $A$ the interaction matrix and $\\vec{\\phi}$ the incoming pertubation on the DM. This equation can be solved in a linear least squares sense,\n",
    "\n",
    "$$\\vec{\\phi} = \\left(A^TA\\right)^{-1} A^T\\vec{S}.$$\n",
    "\n",
    "The matrix $\\left(A^TA\\right)^{-1} A^T$ can be found by applying a pseudo-inverse operation on the matrix $A$. A regularized version of this is implemented in HCIpy with the inverse_tikhonov function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rcond = 1E-3\n",
    "reconstruction_matrix = inverse_tikhonov(slopes.transformation_matrix, rcond=rcond, svd=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The wavefront is then initialized with the Magallen aperture and the sensing wavelength. To have something to measure and correct we put a random shape on the DM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "deformable_mirror.random(0.2 * wavelength_wfs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now lets setup the parameters of our AO system. The first step is to choose an integration time for the exposures. We choose an exposure time of 1 ms, so we are running our AO system at 1 kHz. For the controller we choose to use a leaky integrator which has been proven to be a robust controller. The leaky integrator has two parameters, the leakage and the gain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "delta_t = 1E-3\n",
    "leakage = 0.01\n",
    "gain = 0.5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Initialize our wavefront and setup the propagator for evaluation of the PSF."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wf_wfs = Wavefront(magellan_aperture, wavelength_wfs)\n",
    "wf_wfs.total_power = zero_magnitude_flux * 10**(-stellar_magnitude/2.5) * delta_t\n",
    "print(\"Photon flux per frame {:g}\".format(wf_wfs.total_power))\n",
    "\n",
    "spatial_resolution = wavelength_wfs / telescope_diameter\n",
    "focal_grid = make_focal_grid(q=8, num_airy=20, spatial_resolution=spatial_resolution)\n",
    "\n",
    "propagator = FraunhoferPropagator(pupil_grid, focal_grid)\n",
    "\n",
    "Inorm = propagator.forward(wf_wfs).power.max()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's check the current PSF that is created by the deformed mirror."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PSF_in = propagator.forward(deformable_mirror.forward(wf_wfs)).power / Inorm\n",
    "input_strehl = PSF_in.max()\n",
    "\n",
    "imshow_field(np.log10(PSF_in), vmin=-5, vmax=0)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we are ready to run the system in closed loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_closed_loop_animation():\n",
    "    \n",
    "    wf_ref = Wavefront(magellan_aperture, wavelength_wfs)\n",
    "    PSF = propagator.forward(wf_ref).power\n",
    "    PSF /= PSF.max()\n",
    "    \n",
    "    fig = plt.figure(figsize=(12,4))\n",
    "    plt.subplot(1,3,1)\n",
    "    plt.title('DM surface shape')\n",
    "    im1 = imshow_field(deformable_mirror.surface, vmin=-1e-6, vmax=1e-6, cmap='bwr')\n",
    "        \n",
    "    plt.subplot(1,3,2)\n",
    "    plt.title('Wavefront sensor output')\n",
    "    im2 = imshow_field(image_ref, pwfs.output_grid)\n",
    "    \n",
    "    plt.subplot(1,3,3)\n",
    "    plt.title('Science image plane')\n",
    "    im3 = imshow_field(np.log10(PSF), vmin=-5, vmax=0)\n",
    "    \n",
    "    plt.close(fig)\n",
    "\n",
    "    def animate(t):\n",
    "        wf_dm = deformable_mirror.forward(wf_wfs)\n",
    "        wf_pyr = pwfs.forward(wf_dm)\n",
    "\n",
    "        camera.integrate(wf_pyr, 1)\n",
    "        wfs_image = large_poisson(camera.read_out()).astype(np.float)\n",
    "        wfs_image /= np.sum(wfs_image)\n",
    "\n",
    "        diff_image = wfs_image-image_ref\n",
    "        deformable_mirror.actuators = (1-leakage) * deformable_mirror.actuators - gain * reconstruction_matrix.dot(diff_image)\n",
    "\n",
    "        phase = magellan_aperture * deformable_mirror.surface\n",
    "        phase -= np.mean(phase[magellan_aperture>0])\n",
    "        \n",
    "        psf = propagator.forward(wf_dm).power\n",
    "        im1.set_data(*pupil_grid.separated_coords, (magellan_aperture * deformable_mirror.surface).shaped)\n",
    "        im2.set_data(*pwfs.output_grid.separated_coords, wfs_image.shaped)\n",
    "        im3.set_data(*focal_grid.separated_coords, np.log10(psf.shaped/psf.max()))\n",
    "\n",
    "        return [im1, im2, im3]\n",
    "    \n",
    "    num_time_steps=21\n",
    "    time_steps = np.arange(num_time_steps)\n",
    "    anim = animation.FuncAnimation(fig, animate, time_steps, interval=160, blit=True)\n",
    "    return HTML(anim.to_jshtml(default_mode='loop'))\n",
    "\n",
    "create_closed_loop_animation()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  },
  "level": "intermediate",
  "thumbnail_figure_index": 1
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
